from typing import Dict, List
from llama_index.llms.openai import OpenAI
from llama_index.llms.openai.utils import ALL_AVAILABLE_MODELS, CHAT_MODELS
from env import Env  # 将API_KEY配置在env.py文件里面

# 定义 DeepSeek 模型及其上下文长度
DEEPSEEK_MODELS: Dict[str, int] = {
    "deepseek-ai/DeepSeek-R1-Distill-Llama-8B": 128000,
}

# 更新全局模型列表
ALL_AVAILABLE_MODELS.update(DEEPSEEK_MODELS)
CHAT_MODELS.update(DEEPSEEK_MODELS)

# API 密钥（建议使用环境变量管理）
API_KEY = Env.API_KEY


def get_siliconflow_llm(api_key=API_KEY, model='deepseek-ai/DeepSeek-R1-Distill-Llama-8B',
                        api_base='https://api.siliconflow.cn/v1', temperature=0.6, **kwargs) -> OpenAI:
    """
    创建并返回一个配置好的OpenAI实例，用于与DeepSeek模型交互。

    Args:
        api_key (str): API 密钥，默认为全局 API_KEY。
        model (str): 模型名称，默认为 'deepseek-ai/DeepSeek-R1-Distill-Llama-8B'。
        api_base (str): API 基础 URL，默认为 'https://api.siliconflow.cn/v1'。
        temperature (float): 生成文本的随机性，默认为 0.6。
        **kwargs: 其他传递给OpenAI构造函数的参数。

    Returns:
        OpenAI: 配置好的OpenAI实例。
    """
    llm = OpenAI(api_key=api_key, model=model, api_base=api_base, temperature=temperature, **kwargs)
    return llm


def chat_with_model(llm: OpenAI) -> None:
    """
    与模型进行持续对话，通过不断追加上下文实现多轮对话。

    Args:
        llm (OpenAI): 配置好的OpenAI实例。
    """
    # 初始化对话上下文
    context = []

    print("开始对话！输入 'exit' 退出。")
    while True:
        # 获取用户输入
        user_input = input("你: ")
        if user_input.lower() == 'exit':
            print("退出对话。")
            break

        # 将用户输入添加到上下文
        context.append({"role": "user", "content": user_input})

        try:
            # 将上下文转换为字符串形式（根据模型输入格式调整）
            prompt = "\n".join([f"{msg['role']}: {msg['content']}" for msg in context])

            # 调用模型生成响应
            response = llm.complete(prompt)
            model_reply = str(response)

            # 将模型响应添加到上下文
            context.append({"role": "assistant", "content": model_reply})

            # 打印模型响应
            print(f"模型: {model_reply}")
        except Exception as e:
            print(f"An error occurred: {e}")


if __name__ == '__main__':
    # 获取配置好的OpenAI实例
    llm = get_siliconflow_llm()

    # 开始对话
    chat_with_model(llm)
